import numpy as np

class PSS:
    zc_map = [25,29,34]
    def __init__(self,t):
        ofd = np.exp(-1j*np.pi*self.zc_map[t]/63*np.arange(63)*np.arange(1,64))
        self.fd = np.zeros((128,),dtype=complex)
        self.fd[1:32] = ofd[32:]
        self.fd[97:]  = ofd[:31]
        self.ofd = np.zeros((62,),dtype=complex)
        self.ofd[:31] = ofd[:31]
        self.ofd[31:] = ofd[32:]

        td = np.fft.ifft(self.fd) 
        self.td = np.zeros((137,),dtype=complex)
        self.td[9:] = np.conj(td[:])
        self.td[:9] = np.conj(td[-9:])
        self.tlen = 137

    def xcorr_ft(self,rfd,sset,fc=874.2e6,sr=1.92e6):
        dlen = rfd.shape[0]-100
        fn = dlen//19200
        halffl = 9600
        r = np.zeros((len(sset),halffl))
        for foi in range(len(sset)):
            f = sset[foi]
            drc = 1+f/fc
            dfps = self.td*np.exp(-1j*np.pi*2.*f/sr*np.arange(len(self.td)))
            for i in range(halffl):
                for j in range(2*fn):
                    ppos = int((i+(j*halffl))/drc)
                    x = np.dot(rfd[ppos:ppos+self.tlen],dfps)
                    r[foi,i] += np.abs(x)
        return r
class SSS:
    def __init__(self,id_1,id_2,slot):
        qp = id_1//30
        q  = (id_1+qp*(qp+1)//2)//30
        mp = id_1+q*(q+1)//2
        m0 = mp%31
        m1 = (m0+mp//31+1)%31
        s_td = 1-2*np.array([0,0,0,0,1,0,0,1,0,1,1,0,0,1,1,1,1,1,0,0,0,1,1,0,1,1,1,0,1,0,1])
        c_td = 1-2*np.array([0,0,0,0,1,0,1,0,1,1,1,0,1,1,0,0,0,1,1,1,1,1,0,0,1,1,0,1,0,0,1])
        z_td = 1-2*np.array([0,0,0,0,1,1,1,0,0,1,1,0,1,1,1,1,1,0,1,0,0,0,1,0,0,1,0,1,0,1,1])
        s0_m0 = np.array([ s_td[(i+m0)%31] for i in range(31)])
        s1_m1 = np.array([ s_td[(i+m1)%31] for i in range(31)])
        c0    = np.array([ c_td[(i+id_2)%31] for i in range(31)])
        c1    = np.array([ c_td[(i+id_2+3)%31] for i in range(31)])
        z1_m0 = np.array([ z_td[(i+m0%8)%31] for i in range(31)])
        z1_m1 = np.array([ z_td[(i+m1%8)%31] for i in range(31)])
        if slot==0:
            ssc2 = s1_m1*c1*z1_m0
            ssc1 = s0_m0*c0
        else:
            ssc2 = s0_m0*c1*z1_m1
            ssc1 = s1_m1*c0
        self.ofd = np.zeros(62)
        self.ofd[::2] = ssc1[:]
        self.ofd[1::2] = ssc2[:]
        self.fd = np.zeros((128,),dtype=complex)
        self.fd[1:32] = self.ofd[31:]
        self.fd[97:]  = self.ofd[:31]  